import os
from langchain_community.chat_models import ChatZhipuAI 
from dotenv import load_dotenv, find_dotenv
from langchain_core.prompts import (
    FewShotChatMessagePromptTemplate,
    ChatPromptTemplate
)

_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

llm = ChatZhipuAI(model=model,
                  temperature=0.9,
    )


examples = [
    {
        "中文":"你好。",
        "英文":"Hello.",
        "日文":"こんにちは。"
    }
]

example_prompt = ChatPromptTemplate.from_messages(
    [('human', '{中文}'), ('ai', '英文：{英文}\n日文：{日文}')]
)

few_shot_prompt = FewShotChatMessagePromptTemplate(
    examples=examples, 
    example_prompt=example_prompt,
)

final_prompt = ChatPromptTemplate.from_messages(
    [
        ('system', 
         '你是一名实时翻译，你能将用户输入的消息同步翻译成英文和日文。'),
        few_shot_prompt,
        ('human', '{input}'),
    ]
)

result = llm.invoke(final_prompt.format(input="我是只快乐的小蜜蜂。"))
print(result.content)
